# -*- coding: utf-8 -*-
"""
Created on Thu Oct 10 13:59:03 2013

@author: jkwong
"""

#PBAR_LiningUpRadZspecImages.py

import PBAR_Zspec, PBAR_Cargo
import numpy as np
import os
import matplotlib.pyplot as plt
from scipy import interpolate
from matplotlib import cm
from scipy import ndimage
import copy
import glob

# Set useful plot variables
plotColors = ['r', 'b', 'g', 'm', 'c', 'y', 'k'] * 10
lineStyles = ['-', '-.', ':', '_', '|'] *  10
markerTypes = ['.', 'o', 'v', '^', '<', '>', '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd']
markerTypes = markerTypes * 2
colorList = np.array(['r', 'b', 'g', 'm', 'c', 'y'])

# list of zpsec detector names - Z1 to Z136
ZspecNameList = np.array(['Z%d'%i for i in np.arange(1, 137)])
# list of bad detector numbers starting at 1
badZspecList = np.array([1,2,3,4,5,6,7,8,20,26,31,33,34,38,39,40,44,53,56,62,68,76,80,125,126,127,128,129,130,131,132,133,134,135,136])
# list of good detector numbers starting at 1
goodZspecList = np.array([i for i in np.arange(1, 137) if (i not in badZspecList)])
# list of bad detector names
badZspecNameList = ZspecNameList[(badZspecList-1)]
# list of good detector names
goodZspecNameList = np.array(['Z%d'%i for i in np.arange(1, 137) if (i not in badZspecList)])

badZspecIndices = badZspecList - 1
goodZspecIndices = goodZspecList - 1

badZspecMask = np.zeros(len(ZspecNameList))
badZspecMask[badZspecIndices] = True
badZspecMask = badZspecMask.astype(bool)

goodZspecMask = np.zeros(len(ZspecNameList))
goodZspecMask[goodZspecIndices] = True
goodZspecMask = goodZspecMask.astype(bool)

dataDir = r'C:\Users\jkwong\Documents\Work\PBAR\data3\BasicScans'
plotDir = r'C:\Users\jkwong\Documents\Work\PBAR\data3\BasicScansPlots'
# define parameters for standard width images
numberTimeSlices = 2000
preCargoTimeSlices = 70

acqTime = 1/60.
figureSize = (16, 10)

#  wBest = np.array([-2.2806503 ,  2.26895887,  6.7898522 ])  # old 10/14/2013
        
wBest = np.array([-2.07779891,  1.88095226,  4.65429156])

#################################
## DEFINE FILE NAMES

cargoConfigBaseDirList = glob.glob(os.path.join(dataDir, '*'))

cargoConfigBaseDirList = cargoConfigBaseDirList[36:]

for (dirIndex, cargoConfigBaseDir) in enumerate(cargoConfigBaseDirList):
    subDirList = glob.glob(os.path.join(cargoConfigBaseDir, '*'))
    for (subDirIndex, dataPath) in enumerate(subDirList):
        print(dataPath)

        a, b = os.path.split(dataPath)
        a, c = os.path.split(a)

        plotSaveDir = os.path.join(plotDir, c, b)
        
        # make paths if not present
        if not os.path.exists(os.path.join(plotDir, c)):
            os.mkdir(os.path.join(plotDir, c))

        if not os.path.exists(plotSaveDir):
            os.mkdir(plotSaveDir)
        
        acquisitionTimeList = []
        zspecScanNumberList = []
        cargoScanNumberList = []
        
        fullFilenameList = []
        filenameList = []
        
        filenameCargoList = []
        fullFilenameCargoList = []
        markerfullFilenameCargoList = []
        
        # find the zspec scan number
        temp = glob.glob(os.path.join(dataPath, '*.npy'))
        # if file doesn't exist skip this set completely
        if temp == []:
            print('No npy file found.  Skipping.')
            continue
        a, filename = os.path.split(temp[0])
        zspecScanNumberList.append(filename[0:4])
        filenameList.append(filename)
        fullFilenameList.append(os.path.join(dataPath, filename))

        
        # get the cargo image name
        temp = glob.glob(os.path.join(dataPath, '*.cargoimage'))
        # if file doesn't exist skip this set completely
        if temp == []:
            print('No cargoimage file found.  Skipping.')
            continue
        a, filename = os.path.split(temp[0])
        cargoScanNumberList.append(filename[0:16])
        filenameCargoList.append(filename)    
        fullFilenameCargoList.append(os.path.join(dataPath, filename))
        
        markerfullFilenameCargoList.append(os.path.join(dataPath, filename.replace('cargoimage', 'cargomarker')))
        acquisitionTimeList.append(acqTime)
        
        
        ######################
        ## READ IN DATA
        
        # Read in marker files
        markerList = []
        for (index, f) in enumerate(markerfullFilenameCargoList):
            markerList.append(PBAR_Cargo.ReadCargoMarker(f))
        
        # Read Zspec
        energyList = []
        dat = []
        datStandardWidth = []
        startBinZspecList = []
        for (index, fullFilename) in enumerate(fullFilenameList):
            (a, b) = PBAR_Zspec.ReadZspecBasicScanNumpy(fullFilename)
            energyList.append(a)
            dat.append(b)
        
            startBinZspec = PBAR_Zspec.FindBasicZspecStart(b)
            startBinZspecList.append(startBinZspec)
            
            temp = b[(startBinZspec-preCargoTimeSlices):,:][0:numberTimeSlices,:,:] # this could have length less than numberTimeSlices    
            
           # This ensure that we have standard size images    
            temp2 = np.zeros((numberTimeSlices, b.shape[1], b.shape[2]))
            if temp.shape[0] < numberTimeSlices:
                temp2[0:temp.shape[0],:,:] = temp
            else:
                temp2 = temp
            datStandardWidth.append(temp2)
        del a, b
        
        # Read Cargo
        numberZspecDetectors = dat[0].shape[1]
        datCargo = []
        datCargoStandardWidth = []
        startBinCargoList = []
        
        for (index, f) in enumerate(fullFilenameCargoList):
            (A,bpp,formatt,flag,low1,high1,low2,high2) = PBAR_Cargo.ReadCargoImage(f)
            datCargo.append(A)
        
            # this should do the same thing
            indices1 = 1 + 4 * np.arange(136) 
            indices2 = 2 + 4* np.arange(136)
            temp = (A[:,indices1] + A[:,indices2])/2
            
            startBinCargo = PBAR_Cargo.FindCargoStart(A)
            startBinCargoList.append(startBinCargo)
            
            # cut section in front and limit to numberTimeSlices number of time slices
            temp = temp[(startBinCargo-preCargoTimeSlices):,:][0:numberTimeSlices,:]
            
            # This ensure that we have standard size images
            # excess is filled with zeros
            temp2 = np.zeros((numberTimeSlices, numberZspecDetectors))
            if temp.shape[0] < numberTimeSlices:
                temp2[0:temp.shape[0],:] = temp
            else:
                temp2 = temp[0:numberTimeSlices,:]
                
            datCargoStandardWidth.append(temp2)    
            
        del A,bpp,formatt,flag,low1,high1,low2,high2
        del indices1, indices2, temp, temp2
        
        # Modify the marker files based on found start times
        markerStandardWidthList = []
        for (index, markers) in enumerate(markerList):
            # first copy over
            temp = copy.deepcopy(markers)
            
            # adjust the x values
            offset = (startBinCargo - preCargoTimeSlices)
            for i in np.arange(len(temp)):
                temp[i]['rec_left'] = temp[i]['rec_left'] - offset
                temp[i]['rec_right'] = temp[i]['rec_right'] - offset
                temp[i]['x'] = temp[i]['x'] - offset
                # see if 'left and 'right exist
                if 'left' in temp[i]:
                    temp[i]['left']['x'] = temp[i]['left']['x'] - offset
                if 'right' in temp[i]:
                    temp[i]['right']['x'] = temp[i]['right']['x'] - offset
            # adjust the y values
            for i in np.arange(len(temp)):
                temp[i]['rec_top'] = np.round((temp[i]['rec_top'] - 1.5) / 4.0)
                temp[i]['rec_bottom'] = np.round((temp[i]['rec_bottom'] - 1.5) / 4.0)
                temp[i]['y'] = np.round((temp[i]['y'] - 1.5) / 4.0)
                # see if 'left and 'right exist
                if 'left' in temp[i]:
                    temp[i]['left']['y'] = np.round((temp[i]['left']['y'] - 1.5) / 4.0)
                if 'right' in temp[i]:
                    temp[i]['right']['y'] = np.round((temp[i]['right']['y'] - 1.5) / 4.0)
            markerStandardWidthList.append(copy.copy(temp))
        
        #############################
        ##  CALCULATE FEATURES
        
        discrim = []
        for (index, d) in enumerate(datStandardWidth):
            temp = {}
            energy = energyList[index].astype(float)
            
            temp['count'] = d.sum(2).astype(float)
            
            energyMatrix = np.tile(energy, (d.shape[0], d.shape[1], 1))
            temp['binMean'] = (d * energyMatrix).sum(2) / d.sum(2)
            
            binMeanMatrix = np.tile(temp['binMean'], (d.shape[2],1,1)).transpose((1, 2, 0))
            temp['binSTD'] = np.sqrt( (( (energyMatrix - binMeanMatrix)**2) * d.astype(float) ).sum(2) / d.sum(2).astype(float)  )
        
            temp['binSTD_binMean'] = temp['binSTD'] / temp['binMean']
            
            # find the bin that is just above 10, 20
            highlowIndex10 = np.argmin( abs(energy - 10) )
            highlowIndex20 = np.argmin( abs(energy - 20) )
            
            temp['multibin_0_10'] = d[:,:,0:highlowIndex10].sum(axis=2)  # bins 0 to 9, so first ten bins
            temp['multibin_10_256'] = d[:,:,highlowIndex10:].sum(axis=2)
            temp['multibin_0_20'] = d[:,:,0:highlowIndex20].sum(axis=2)
            temp['multibin_20_256'] = d[:,:,highlowIndex20:].sum(axis=2)
        
            temp['multibin_10_ratio'] = temp['multibin_0_10'] / temp['multibin_10_256']
            temp['multibin_20_ratio'] = temp['multibin_0_20'] / temp['multibin_20_256']
        
            # QUANTILES
            datTempCumSum = d.cumsum(axis = 2).astype(float) # cumulative sum
            datTempCumSum = datTempCumSum / np.tile(datTempCumSum[:,:,-1], (32, 1, 1)).transpose((1, 2, 0)) # normalize the cumulative sum
            qIntervals = [10., 25., 50., 75., 90.]
            for j in range(len(qIntervals)):
                statName = 'q_%02d' % qIntervals[j]# apparently needs a temporary variable
                temp[statName] = energy[np.argmin(abs(datTempCumSum - (qIntervals[j]/100)), axis = 2)]
            temp['q_range'] = temp['q_75']- temp['q_25']
            temp['q_range_ratio'] = temp['q_range'] / temp['q_50']
            temp['binMean_q50'] = temp['binMean'] - temp['q_50']
            discrim.append(temp)
        
        # Calculate the dist, linear combination

        for index in np.arange(len(discrim)):
            # make features array (num features X time slices X number of detectors)
            featuresMatrix = np.zeros((3, discrim[index]['binMean'].shape[0], discrim[index]['binMean'].shape[1]))
#            featuresMatrix[0,:,:] = discrim[index]['binMean'] # old 10/14/2013
#            featuresMatrix[1,:,:] = discrim[index]['binSTD']
#            featuresMatrix[2,:,:] = discrim[index]['binMean_q50']

            featuresMatrix[0,:,:] = discrim[index]['binMean']
            featuresMatrix[1,:,:] = discrim[index]['binSTD']
            featuresMatrix[2,:,:] = discrim[index]['multibin_20_ratio']
            
            dist0 = featuresMatrix[0,:,:] *wBest[0] + featuresMatrix[1,:,:] *wBest[1] + featuresMatrix[2,:,:] *wBest[2]
            # clean up bad values
            cutt = np.isnan(dist0)
            dist0[cutt] = 0.0
            
            discrim[index]['dist0'] = dist0
            
        # Cut in radiography image of standard size
        threshold = 1.5e7 # only consider parts of the zspec image that have rad below this value
        cargoZspecStandardWidthMask = []
        for i in np.arange(len(datCargoStandardWidth)):
            cargoZspecStandardWidthMask.append(datCargoStandardWidth[i] <= threshold)
        
        # Clean up radiograph mask
        cargoZspecStandardWidthCleanedMask = []
        for i in np.arange(len(datCargoStandardWidth)):
            temp = cargoZspecStandardWidthMask[i]
            # do some stuff
            temp = ndimage.binary_opening(temp)
            cargoZspecStandardWidthCleanedMask.append(temp)

        
        ###################
        ## PLOTS 
        index = 0
        
        ## Unadjusted Images ##
        
        ## Plot raw unadjusted images, both
        # for zspec plot the total counts

        logIntensity = True
        plotMarkers = True
        
        plt.figure(figsize = figureSize)
        plt.grid()
        
        if logIntensity:
            intensity = np.log(datCargo[index].T)
        else:
            intensity = datCargo[index].T
        
        plt.imshow(intensity, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
        plt.colorbar()
        plt.title(filenameCargoList[index] + ' Original Size')
        plt.xlim((0, 2400))
        
        if plotMarkers:
            marker = markerList[index]
            for i, mark in enumerate( marker):
                x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
                y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
                plt.plot(x, y, 'r')
                plt.plot(mark['x'], mark['y'], 'xr')
                if 'left' in mark:
                    plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
                if 'right' in mark:
                    plt.plot(mark['right']['x'], mark['right']['y'], 'go')
 
       # save image to file
        plotFullFilename = os.path.join(plotSaveDir, '%s_OriginalSize.png' %filenameCargoList[index])
        plt.savefig(plotFullFilename)
        plt.close('all')
                    
        plt.figure(figsize = figureSize)
        plt.grid()
        
        
        # Zspec, original size
        
        if logIntensity:
            intensity = np.log(dat[index].sum(axis = 2).T)
        else:
            intensity = dat[index].sum(axis = 2).T
        
        plt.imshow(intensity, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
        plt.colorbar()
        plt.title(filenameList[index] + ' Original Size')
        plt.xlim((0, 2400))

        # save image to file
        plotFullFilename = os.path.join(plotSaveDir, '%s_OriginalSize.png' %filenameList[index])
        plt.savefig(plotFullFilename)
        plt.close('all')
        
        ## Standard Size Images ##
        
        ## Cargo standard size with markers
        # for zspec plot the total counts
        logIntensity = 1
        plotMarkers = True
        plt.figure(figsize = figureSize)
        plt.grid()
        
        if logIntensity:
            intensity = np.log(datCargoStandardWidth[index].T)
        else:
            intensity = datCargoStandardWidth[index].T
            
        plt.imshow(intensity, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
        plt.colorbar()
        plt.title(filenameCargoList[index] + ', Standard Size')
        plt.xlim((0, 2400))
        
        if plotMarkers:
            marker = markerStandardWidthList[index]
            for i, mark in enumerate( marker):
                x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
                y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
                plt.plot(x, y, 'r')
                plt.plot(mark['x'], mark['y'], 'xr')
                if 'left' in mark:
                    plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
                if 'right' in mark:
                    plt.plot(mark['right']['x'], mark['right']['y'], 'go')
        
        # save image to file
        plotFullFilename = os.path.join(plotSaveDir, '%s_StandardSize.png' %filenameCargoList[index])
        plt.savefig(plotFullFilename)        
        plt.close('all')
        
        # Zspec image, counts with markers        
        plt.figure(figsize = figureSize)
        plt.grid()
        featureName1 = 'count'
        t = plt.imshow(discrim[index][featureName1].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
        plt.colorbar()
        if plotMarkers:
            marker = markerStandardWidthList[index]
            for i, mark in enumerate( marker):
                x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
                y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
                plt.plot(x, y, 'r')
                plt.plot(mark['x'], mark['y'], 'xr')
                if 'left' in mark:
                    plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
                if 'right' in mark:
                    plt.plot(mark['right']['x'], mark['right']['y'], 'go')
        plt.title(filenameList[index] + ', Standard Size')
        plt.xlim((0, 2400))
                
                
        # save image to file
        plotFullFilename = os.path.join(plotSaveDir, '%s_%s_StandardSize.png' %(filenameList[index], featureName1))
        plt.savefig(plotFullFilename)        

        plt.close('all')
        
        
        # Zspec dist image
        plt.figure(figsize = figureSize)
        plt.grid()
        featureName1 = 'dist0'
        t = plt.imshow(discrim[index][featureName1].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
        if plotMarkers:
            marker = markerStandardWidthList[index]
            for i, mark in enumerate( marker):
                x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
                y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
                plt.plot(x, y, 'r')
                plt.plot(mark['x'], mark['y'], 'xr')
                if 'left' in mark:
                    plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
                if 'right' in mark:
                    plt.plot(mark['right']['x'], mark['right']['y'], 'go')        
        plt.colorbar()
        plt.title('%s, %s ' %(filenameList[index], featureName1))
        plt.xlim((0, 2200))
        
        plotFullFilename = os.path.join(plotSaveDir, '%s_%s_StandardSize.png' %(filenameList[index], featureName1))
        plt.savefig(plotFullFilename)   
        
        # Zspec dist image with mask
        plt.figure(figsize = figureSize)
        plt.grid()
        featureName1 = 'dist0'
        im = copy.copy(discrim[index][featureName1].T)
        minValue = im.min()
        im[~cargoZspecStandardWidthCleanedMask[index].T] = minValue
        t = plt.imshow(im, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
        plt.colorbar()
        plt.title('Linear Combination')
        plt.xlim((0, 2200))
        plt.show()

        plotFullFilename = os.path.join(plotSaveDir, '%s_%s_Mask_StandardSize.png' %(filenameList[index], featureName1))
        plt.savefig(plotFullFilename)  
        
        
        plt.close('all')
        ### Scatter Plots ###
        
        #  All data points and mask, and those in the boxes
        plt.figure(figsize = figureSize)
        plt.grid()
        index = 0
        feature1 = 'count'        
        feature2 = 'dist0'
        ## SCATTER - dist0 vs count
        
        # x, y for the whole image
        x1 = discrim[index][feature1][:,goodZspecIndices].flatten()
        y1 = discrim[index][feature2][:,goodZspecIndices].flatten()
        
        # reduce the mask to good detectors
        cut = cargoZspecStandardWidthCleanedMask[index][:, goodZspecIndices]
        # x values in mask, goog detectors        
        x2 = discrim[index][feature1][:,goodZspecIndices][cut].flatten()
        # y values in mask, goog detectors        
        y2 = discrim[index][feature2][:,goodZspecIndices][cut].flatten()
        
        plt.plot(x1, y1, '.k', alpha = 0.15, markersize = 10, label = 'All Points')
        plt.plot(x2, y2, '.b', alpha = 0.15, markersize = 10, label = 'In Mask')
        
        # plot the points in the markers        
        marker = markerStandardWidthList[index]
        for i, mark in enumerate(marker):
            x_range = np.array((mark['rec_left'], mark['rec_right'])) # left < right
            y_range = np.array((mark['rec_top'], mark['rec_bottom'])) # top  < bottom
            xarray = np.arange(discrim[index][feature1].shape[0])
            yarray = np.arange(discrim[index][feature1].shape[1])
            xcut = (xarray > x_range[0]) & (xarray < x_range[1])
            ycut = (yarray > y_range[0]) & (yarray < y_range[1])
            
            x = discrim[index][feature1][:,goodZspecMask & ycut][xcut,:].flatten()
            y = discrim[index][feature2][:,goodZspecMask & ycut][xcut,:].flatten()
            
            plt.plot(x, y, '.', markersize = 15, alpha  = 0.3, marker = markerTypes[i/6], color = colorList[i%6], label = mark['target'])
        
        plt.title(filenameList[index] + ', Linear Combination Vs Counts')
        plt.legend()
        plt.xlabel(feature1)
        plt.ylabel(feature2)        
        plt.show()

        plotFullFilename = os.path.join(plotSaveDir, '%s_%s_%s_AllPoints_Mask_Markers.png' %(filenameList[index], featureName1, feature2))
        plt.savefig(plotFullFilename)  
        
        plt.close('all')
        

        #  data points in mask, and those in mask and box 
        plt.figure(figsize = figureSize)
        plt.grid()
        feature1 = 'count'        
        feature2 = 'dist0'        
        ## SCATTER - dist0 vs count
        
        # reduce the mask to good detectors
        cut = cargoZspecStandardWidthCleanedMask[index][:, goodZspecIndices]
        # x values in mask, goog detectors        
        x2 = discrim[index][feature1][:,goodZspecIndices][cut].flatten()
        # y values in mask, goog detectors        
        y2 = discrim[index][feature2][:,goodZspecIndices][cut].flatten()
        
        plt.plot(x2, y2, '.k', alpha = 0.15, markersize = 10, label = 'In Mask')
        
        # plot the points in the markers        
        marker = markerStandardWidthList[index]
        for i, mark in enumerate(marker):
            x_range = np.array((mark['rec_left'], mark['rec_right'])) # left < right
            y_range = np.array((mark['rec_top'], mark['rec_bottom'])) # top  < bottom
            xarray = np.arange(discrim[index][feature1].shape[0])
            yarray = np.arange(discrim[index][feature1].shape[1])
            xcut = (xarray > x_range[0]) & (xarray < x_range[1])
            ycut = (yarray > y_range[0]) & (yarray < y_range[1])
            
            x = discrim[index][feature1][:,goodZspecMask & ycut][xcut,:].flatten()
            y = discrim[index][feature2][:,goodZspecMask & ycut][xcut,:].flatten()
            
            plt.plot(x, y, '.', markersize = 15, alpha  = 0.3, marker = markerTypes[i/6], color = colorList[i%6], label = mark['target'])
        
        plt.title(filenameList[index] + ', Linear Combination Vs Counts, Special Points Not Necessarily in Mask')
        plt.legend()
        plt.xlabel(feature1)
        plt.ylabel(feature2)
        plt.show()
        
        plotFullFilename = os.path.join(plotSaveDir, '%s_%s_%s_Mask_Markers.png' %(filenameList[index], featureName1, feature2))
        plt.savefig(plotFullFilename)  

        plt.close('all')
